import torch
import matplotlib.pyplot as plt
from utils.set_datasets import init_datasets
from models.select_model import select_model
import torch as t
from torch import nn, optim
from utils.train_and_test import train_and_test
from utils.config import *
from utils.log_utils import *
from utils.send_email_utils import *
from pruner import SNIP

def apply_prune_mask(model, keep_masks):
    prunable_layer = filter(
        lambda layer: isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear), model.modules())

    for layer, keep_masks in zip(prunable_layer, keep_masks):
        assert (layer.weight.shape == keep_masks.shape)
        def hook_factory(keep_masks):

            def hook(grads):
                return grads * keep_masks

            return hook

        layer.weight.data[keep_masks == 0.] = 0.
        layer.weight.register_hook(hook_factory(keep_masks))

def cal_sum_each_layer_weight(model):
    each_layer = dict()
    for layer in model.modules():
        if isinstance(layer, (nn.Conv2d, nn.Linear)):
            each_layer[layer] = t.sum(layer.weight.data)
    return each_layer

def cal_sum_each_layer_weight_ratio(keep_masks):
    ratio = []
    total = 0
    remained = 0
    for mask in keep_masks:
        ratio.append((mask.sum()/mask.numel()).item()*100)
        total += mask.numel()
        remained += mask.sum()
    print("整个网络的修剪比率是:", (remained/total * 100))
    return ratio
def plot_ration(ration):
    x = list(i for i in range(1, 20))
    plt.figure(figsize=(16, 8))
    plt.plot(x, ration, marker='o', markersize=3,  color='red', alpha=0.8)
    for xi, yi in zip(x, ration):
        plt.annotate(
            f"({yi:.2f})",  # 显示的文本内容
            xy=(xi, yi),  # 箭头指向的坐标
            xytext=(xi+0.1, yi + 0.1),  # 文本的坐标
            arrowprops=None, # 箭头的样式
            fontsize=12,
            color='red'
        )

    plt.title('snip prune each layer ratio')
    plt.xlabel('layer')
    plt.ylabel('each layer ratio')
    plt.show()
def plot_weight(sum_weight1, sum_weight2):
    x = list(i for i in range(1, 20))
    y1 = list(float(i.cpu().detach().numpy()) for i in sum_weight1.values())
    y2 = list(float(i.cpu().detach().numpy()) for i in sum_weight2.values())
    plt.plot(x, y1, marker='o', markersize=3, color='red')
    plt.plot(x, y2, marker='o', markersize=3, color='blue')
    plt.legend()
    plt.title('SNIP after and before prune weight')
    plt.xlabel('layer')
    plt.ylabel('weight')
    plt.show()
def cal_sum_weight(model):
    sum_weight = dict()
    for layer in model.modules():
        if isinstance(layer, (nn.Conv2d, nn.Linear)):
            sum_weight[layer] = t.sum(t.sum(t.abs(layer.weight.data), dim=0))
    return sum_weight


def train(config):
    trainloader, testloader = init_datasets(dataset_name=config.dataset, batch_size=config.batch_size)
    model = select_model(config.network)
    device = t.device('cuda' if t.cuda.is_available() else 'cpu')
    model.to(device)
    sum_weight1 = cal_sum_weight(model)
    keep_masks = SNIP.snip(model, keep_ratios=0.05, trainloader=trainloader, devive=device)
    apply_prune_mask(model, keep_masks)
    ration = cal_sum_each_layer_weight_ratio(keep_masks)
    sum_weight2 = cal_sum_weight(model)

    plot_weight(sum_weight1, sum_weight2)


    plot_ration(ration)








    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=config.l2_regularization)
    train_best_accuracy, train_best_epoch, test_best_accuracy, test_best_epoch = train_and_test(config, model, device,
                                                                                                criterion, optimizer,
                                                                                                trainloader, testloader)
    tensorboard_read(config)
    print_model_param(model)

    content = mail_txt(train_best_accuracy, train_best_epoch, test_best_accuracy, test_best_epoch, config)
    send_email("1659078644@qq.com", content, config)
    zip_my_log(config)




if __name__ == '__main__':
    config = init_config()
    train(config)








    
